/**
 * 
 */
package edu.umd.clip.lm.model.decoding;

import java.util.*;

import edu.umd.clip.lm.factors.*;
import edu.umd.clip.lm.util.*;
import edu.umd.clip.lm.util.tree.BinaryTree;
import edu.umd.clip.lm.model.*;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class TrellisSlice {
	//private List<SliceItem> elements; 
	protected HashMap<TrellisItem,ArrayList<TrellisItem>> curr2prev;
	protected HashMap<TrellisItem,ArrayList<TrellisItem>> prev2curr;
	protected HashMap<ClusterId, TrellisItem> itemsByCluster;
	protected double bestProb = 0.0;
	long overtFactors;
	BinaryPrefix wordPrefix;
	private Trellis trellis;
	private int time;
	private double totalProb = Double.NaN;
	private boolean dummy = false; // whether the slice is just used to prefix p(w t | ctx) 
	private FuturePredictions<TrellisItem> predictions;
	private double sliceLogScale;
	private double cumulativeLogScale;
	
	/**
	 * @param overtFactors
	 * @param trellis
	 * @param time
	 */
	public TrellisSlice(long overtFactors, Trellis trellis, int time) {
		this.overtFactors = overtFactors;
		Experiment experiment = Experiment.getInstance();
		int value = FactorTuple.getValue(overtFactors, experiment.getTupleDescription().getMainFactorIndex());
		wordPrefix = experiment.getWordPrefix(value);

		this.trellis = trellis;
		this.time = time;
		//this.elements = new LinkedList<SliceItem>();
		this.prev2curr = new HashMap<TrellisItem,ArrayList<TrellisItem>>(100);
		this.curr2prev = new HashMap<TrellisItem,ArrayList<TrellisItem>>(100);
		this.itemsByCluster = new HashMap<ClusterId, TrellisItem>(200);
		predictions = new FuturePredictions<TrellisItem>();
	}
	
	protected void insertItem(TrellisItem item) {
		//elements.add(item);
		assert(item.slice == null || item.slice == this);
		item.slice = this;
		//item.tree.setTuple(overtFactors);
		if (trellis.getDebug() > 1) {
			System.out.printf("Slice #%d: insertItem %s, prevProb = %e, treeProb = %e\n", time, 
				item, item.getPrevProb(), item.tree.getTotalProb());
		}
		totalProb = Double.NaN;
	}

	protected double computeTotalProb() {
		double prob = 0.0;
		//System.err.println("totalProb, slice #" + Integer.toString(time));
		for(TrellisItem item : getElements()) {
			double elementProb = item.getPrevProb() * item.tree.getTotalProb(); 
			//System.err.println("  " + Double.toString(elementProb));
			prob += elementProb;
		}
		return prob;
	}
	
	protected void changeForwardProbability(TrellisItem item, double probChange) {
		TrellisSlice nextSlice = getNext();
		if (nextSlice == null) return;
		Collection<TrellisItem> nextItems = nextSlice.prev2curr.get(item);
		if (nextItems != null) {
			TrellisSlice nextNextSlice = nextSlice.getNext();
			for(TrellisItem i : nextItems) {
				i.addPrevProb(probChange);
				if (nextNextSlice != null) {
					// TODO: this can get ugly in higher order models
					nextSlice.changeForwardProbability(i, probChange * i.tree.getTotalProb());
				}
			}
		}
	}
	
	@SuppressWarnings("null")
	protected void removeNodes(HashSet<TrellisItem> nodes) {
		TrellisSlice prevSlice = getPrevious();
		HashSet<TrellisItem> prevItems = null;
		if (prevSlice != null) {
			prevItems = new HashSet<TrellisItem>();
		}
		
		TrellisSlice nextSlice = getNext();
		HashSet<TrellisItem> nextItems = null;
		HashSet<TrellisItem> nextToRemove = null;
		if (nextSlice != null) {
			nextItems = new HashSet<TrellisItem>();
			nextToRemove = new HashSet<TrellisItem>();
		}
		
		for(TrellisItem item : nodes) {
			if (prevSlice != null) {
				List<TrellisItem> prevs = curr2prev.get(item);
				if (prevs != null) {
					for(TrellisItem i : prevs) {
						if (!prevItems.contains(i)) {
							List<TrellisItem> lst = prev2curr.get(i);
							lst.removeAll(nodes);
							prevItems.add(i);
						}
					}
					curr2prev.remove(item);
				}
			}
			
			if (nextSlice != null) {
				List<TrellisItem> nexts = nextSlice.prev2curr.get(item);
				if (nexts != null) {
					for(TrellisItem i : nexts) {
						if (!nextItems.contains(i)) {
							List<TrellisItem> lst = nextSlice.curr2prev.get(i);
							lst.removeAll(nodes);
							if (lst.isEmpty()) {
								nextToRemove.add(i);
							}
							nextItems.add(i);
						}
					}
					nextSlice.prev2curr.remove(item);
				}
			}
		}
		//elements.removeAll(nodes);
		if (nextToRemove != null && nextToRemove.size() > 0) {
			nextSlice.removeNodes(nextToRemove);
		}
	}
	
	public double getTotalProb() {
		if (Double.isNaN(totalProb)) {
			totalProb = computeTotalProb();
		}
		return totalProb;
	}
	/**
	 * @return the time
	 */
	public int getTime() {
		return time;
	}
	
	public TrellisSlice getPrevious() {
		return trellis.getPreviousSlice(this);
	}
	
	public TrellisSlice getNext() {
		return trellis.getNextSlice(this);
	}

	/**
	 * @return the elements
	 */
	public Set<TrellisItem> getElements() {
		return curr2prev.keySet();
	}
/*
	protected void setElements(List<SliceItem> elements) {
		this.elements = elements;
	}
*/	
	/**
	 * @return the overtFactors
	 */
	public long getOvertFactors() {
		return overtFactors;
	}

	public ArrayList<TrellisItem> getPreviousItems(TrellisItem item) {
		return curr2prev.get(item);
	}
	
	protected TrellisItem findSameItemInCluster(ClusterId clusterid) {
		TrellisItem item = itemsByCluster.get(clusterid);
		return item;
	}
	
	protected void addItemToCluster(ClusterId clusterid, TrellisItem item) {
		itemsByCluster.put(clusterid, item);
		if (trellis.getDebug() > 1) {
			System.out.printf("addItemToCluster: item %s, cluster: (%s)\n", 
					item.toString(), clusterid);
		}
	}
	
	protected boolean checkPrevProb(TrellisItem item) {
		ArrayList<TrellisItem> prevLinks = curr2prev.get(item);
		if (prevLinks == null) return true;
		
		double totalPrev = 0;
		for(TrellisItem prev : prevLinks) {
			totalPrev += prev.getPrevProb() * prev.tree.getTotalProb();
		}
		if (!ProbMath.approxEqual(totalPrev, item.getPrevProb())) {
			System.out.printf("%s has wrong prev prob: %e, should be %e\n", item, item.getPrevProb(), totalPrev);
			System.out.print("prev links:");
			for(TrellisItem prev : prevLinks) {
				System.out.printf(" %s", prev);
			}
			System.out.println();
			return false;
		}
		return true;
	}
	
	public boolean checkLinks() {
		if (trellis.getDebug() == 0) return true;
		
		TrellisSlice prevSlice = getPrevious();
		for(TrellisItem item : getElements()) {
			ArrayList<TrellisItem> prevLinks = curr2prev.get(item);
			if (item == null) {
				System.err.println("must have links from previous slice");
				return false;
			}
			for(TrellisItem prevItem : prevLinks) {
				if (!prevSlice.getElements().contains(prevItem)) {
					System.err.println("item is not present in the previous slice");
					return false;
				}
				if (!prev2curr.get(prevItem).contains(item)){
					System.err.println("backreference s not present");
					return false;
				}
			}
			
		}
		return true;
	}
	
	public int numItems() {
		return curr2prev.size();
	}

	/**
	 * @return the trellis
	 */
	public Trellis getTrellis() {
		return trellis;
	}

	protected boolean isDummy() {
		return dummy;
	}

	protected void setDummy(boolean dummy) {
		this.dummy = dummy;
	}
	
	protected void makeGhost() {
		for(TrellisItem item : getElements()) {
			item.makeGhost();
		}
	}

	public void addPrediction(BinaryTree<HistoryTreePayload>[] nodes, TrellisItem prevItem) {
		predictions.addPrediction(nodes, prevItem);
	}

	public FuturePredictions<TrellisItem> getPredictions() {
		return predictions;
	}
	
	public void clearPredictions() {
		predictions = null;
	}
	
	@SuppressWarnings("unchecked")
	protected void split(TrellisItem item, TrellisItem cutOff) {
		assert(item.slice == this);
		cutOff.slice = this;
		// copy links from the previous slice
		{
			ArrayList<TrellisItem> prevLinks = curr2prev.get(item);
			if (prevLinks != null) {
				curr2prev.put(cutOff, (ArrayList<TrellisItem>)prevLinks.clone());
				for(TrellisItem i : prevLinks) {
					prev2curr.get(i).add(cutOff);
				}
			}
		}
		// copy links to the following slice if it exists
		{
			TrellisSlice nextSlice = getNext();
			if (nextSlice != null) {
				ArrayList<TrellisItem> fwdLinks = nextSlice.prev2curr.get(item);
				if (fwdLinks != null) {
					nextSlice.prev2curr.put(cutOff, (ArrayList<TrellisItem>)fwdLinks.clone());
					for(TrellisItem i : fwdLinks) {
						nextSlice.curr2prev.get(i).add(cutOff);
					}
				}
			}
		}
		insertItem(cutOff);
	}
	
	protected void updateForwardProbabilities(TrellisItem item, double probDiff) {
		assert(item.slice == this);
		TrellisSlice nextSlice = getNext();
		if (nextSlice != null) {
			Collection<TrellisItem> prevs = nextSlice.prev2curr.get(item);
			if (prevs != null) {
				for(TrellisItem i : prevs) {
					i.addPrevProb(probDiff);
				}
			}
		}
	}

	public double getSliceLogScale() {
		return sliceLogScale;
	}

	public double getCumulativeLogScale() {
		return cumulativeLogScale;
	}
	
	protected void scale(double scale) {
		for(TrellisItem item : getElements()) {
			item.scale(scale);
		}
		this.sliceLogScale -= Math.log10(scale);
		
		this.cumulativeLogScale = time > 2 ? getPrevious().cumulativeLogScale : 0.0;
		this.cumulativeLogScale += sliceLogScale;
	}
	
	protected void setCumulativeLogScale(double logScale) {
		cumulativeLogScale = logScale;
	}
	
	protected double getMaxScale() {
		double maxScale = 0.0;
		for(TrellisItem item : getElements()) {
			double scale = item.tree.getProbTree().getScale();
			if (scale > maxScale) {
				maxScale = scale;
			}
		}
		return maxScale > 0.0 ? maxScale : 1.0;
	}
}
